from flask import Flask, render_template, jsonify, request, Response, send_file
import ipaddress
import socket
import threading
from ping3 import ping
import time
import requests
from requests.exceptions import RequestException
import re
from concurrent.futures import ThreadPoolExecutor, as_completed
import queue
import urllib3
import warnings
import base64
import csv
from io import StringIO
import datetime
import json
import sqlite3
import os
from apscheduler.schedulers.background import BackgroundScheduler
from apscheduler.triggers.cron import CronTrigger
from apscheduler.triggers.interval import IntervalTrigger

# 禁用SSL警告
urllib3.disable_warnings(urllib3.exceptions.InsecureRequestWarning)
warnings.filterwarnings("ignore")

app = Flask(__name__)
# 创建一个线程安全的队列来存储结果
scan_results = queue.Queue()
# 设置线程池最大线程数
MAX_WORKERS = 100
# 创建调度器
scheduler = BackgroundScheduler()
scheduler.start()

# 数据库初始化
def init_db():
    db_path = os.path.join(os.path.dirname(__file__), 'scanner.db')
    conn = sqlite3.connect(db_path)
    c = conn.cursor()
    c.execute('''
        CREATE TABLE IF NOT EXISTS scan_history (
            id INTEGER PRIMARY KEY AUTOINCREMENT,
            scan_time TIMESTAMP,
            network TEXT,
            results TEXT
        )
    ''')
    c.execute('''
        CREATE TABLE IF NOT EXISTS device_info (
            ip TEXT PRIMARY KEY,
            last_seen TIMESTAMP,
            device_type TEXT,
            model TEXT,
            firmware TEXT,
            ports TEXT,
            auth_info TEXT
        )
    ''')
    c.execute('''
        CREATE TABLE IF NOT EXISTS scheduled_tasks (
            id INTEGER PRIMARY KEY AUTOINCREMENT,
            name TEXT,
            network TEXT,
            schedule_type TEXT,
            schedule_value TEXT,
            last_run TIMESTAMP,
            enabled INTEGER DEFAULT 1
        )
    ''')
    c.execute('''
        CREATE TABLE IF NOT EXISTS device_status (
            ip TEXT PRIMARY KEY,
            last_check TIMESTAMP,
            status TEXT,
            response_time REAL,
            error_count INTEGER DEFAULT 0,
            last_error TEXT
        )
    ''')
    conn.commit()
    conn.close()

def load_scheduled_tasks():
    """加载已保存的定时任务"""
    db_path = os.path.join(os.path.dirname(__file__), 'scanner.db')
    conn = sqlite3.connect(db_path)
    c = conn.cursor()
    c.execute('SELECT * FROM scheduled_tasks WHERE enabled = 1')
    tasks = c.fetchall()
    conn.close()

    for task in tasks:
        task_id, name, network, schedule_type, schedule_value, last_run, enabled = task
        if schedule_type == 'interval':
            # 间隔类型的任务（每X分钟）
            minutes = int(schedule_value)
            scheduler.add_job(
                func=scheduled_scan,
                trigger=IntervalTrigger(minutes=minutes),
                args=[network, name],
                id=f'task_{task_id}',
                name=name,
                replace_existing=True
            )
        elif schedule_type == 'cron':
            # Cron类型的任务（指定时间）
            scheduler.add_job(
                func=scheduled_scan,
                trigger=CronTrigger.from_crontab(schedule_value),
                args=[network, name],
                id=f'task_{task_id}',
                name=name,
                replace_existing=True
            )

def scheduled_scan(network, task_name):
    """执行定时扫描任务"""
    try:
        with app.app_context():
            results = []
            with ThreadPoolExecutor(max_workers=MAX_WORKERS) as executor:
                ip_network = ipaddress.ip_network(network)
                futures = [executor.submit(scan_single_ip, str(ip)) for ip in ip_network.hosts()]
                
                for future in as_completed(futures):
                    try:
                        result = future.result()
                        if result:
                            results.append(result)
                    except Exception as e:
                        print(f"Error in scheduled scan: {e}")

            # 保存扫描结果
            save_scan_results(network, results)
            
            # 更新设备状态
            update_device_status(results)
    except Exception as e:
        print(f"Scheduled scan failed: {e}")

def update_device_status(scan_results):
    """更新设备状态"""
    db_path = os.path.join(os.path.dirname(__file__), 'scanner.db')
    conn = sqlite3.connect(db_path)
    c = conn.cursor()
    
    for result in scan_results:
        # 检查设备是否已存在
        c.execute('SELECT error_count FROM device_status WHERE ip = ?', (result['ip'],))
        existing = c.fetchone()
        
        if existing:
            # 重置错误计数
            c.execute('''
                UPDATE device_status 
                SET last_check = ?, status = ?, response_time = ?, error_count = 0, last_error = NULL
                WHERE ip = ?
            ''', (
                datetime.datetime.now(),
                'online',
                result['response_time'],
                result['ip']
            ))
        else:
            # 添加新设备
            c.execute('''
                INSERT INTO device_status (ip, last_check, status, response_time)
                VALUES (?, ?, ?, ?)
            ''', (
                result['ip'],
                datetime.datetime.now(),
                'online',
                result['response_time']
            ))
    
    conn.commit()
    conn.close()

def check_device_status(ip):
    """检查单个设备状态"""
    try:
        response_time = ping(ip, timeout=1)
        if response_time is not None:
            status = 'online'
            error = None
            error_count = 0
        else:
            status = 'offline'
            error = 'Ping timeout'
            error_count = 1
    except Exception as e:
        status = 'error'
        error = str(e)
        error_count = 1
        response_time = None

    # 更新数据库
    db_path = os.path.join(os.path.dirname(__file__), 'scanner.db')
    conn = sqlite3.connect(db_path)
    c = conn.cursor()
    
    c.execute('''
        INSERT OR REPLACE INTO device_status 
        (ip, last_check, status, response_time, error_count, last_error)
        VALUES (?, ?, ?, ?, ?, ?)
    ''', (
        ip,
        datetime.datetime.now(),
        status,
        response_time,
        error_count,
        error
    ))
    
    conn.commit()
    conn.close()
    
    return {
        'ip': ip,
        'status': status,
        'response_time': response_time,
        'error': error
    }

@app.route('/add_schedule', methods=['POST'])
def add_schedule():
    """添加定时任务"""
    try:
        data = request.json
        name = data.get('name')
        network = data.get('network')
        schedule_type = data.get('type')  # 'interval' 或 'cron'
        schedule_value = data.get('value')  # 分钟数或cron表达式
        
        if not all([name, network, schedule_type, schedule_value]):
            return jsonify({'status': 'error', 'message': '缺少必要参数'})
        
        # 保存到数据库
        db_path = os.path.join(os.path.dirname(__file__), 'scanner.db')
        conn = sqlite3.connect(db_path)
        c = conn.cursor()
        c.execute('''
            INSERT INTO scheduled_tasks (name, network, schedule_type, schedule_value)
            VALUES (?, ?, ?, ?)
        ''', (name, network, schedule_type, schedule_value))
        task_id = c.lastrowid
        conn.commit()
        conn.close()
        
        # 添加到调度器
        if schedule_type == 'interval':
            minutes = int(schedule_value)
            scheduler.add_job(
                func=scheduled_scan,
                trigger=IntervalTrigger(minutes=minutes),
                args=[network, name],
                id=f'task_{task_id}',
                name=name
            )
        elif schedule_type == 'cron':
            scheduler.add_job(
                func=scheduled_scan,
                trigger=CronTrigger.from_crontab(schedule_value),
                args=[network, name],
                id=f'task_{task_id}',
                name=name
            )
        
        return jsonify({'status': 'success', 'task_id': task_id})
    except Exception as e:
        return jsonify({'status': 'error', 'message': str(e)})

@app.route('/remove_schedule/<int:task_id>', methods=['POST'])
def remove_schedule(task_id):
    """删除定时任务"""
    try:
        # 从数据库中禁用任务
        db_path = os.path.join(os.path.dirname(__file__), 'scanner.db')
        conn = sqlite3.connect(db_path)
        c = conn.cursor()
        c.execute('UPDATE scheduled_tasks SET enabled = 0 WHERE id = ?', (task_id,))
        conn.commit()
        conn.close()
        
        # 从调度器中移除任务
        scheduler.remove_job(f'task_{task_id}')
        
        return jsonify({'status': 'success'})
    except Exception as e:
        return jsonify({'status': 'error', 'message': str(e)})

@app.route('/api/schedules', methods=['GET'])
def list_schedules():
    """获取所有定时任务"""
    try:
        db_path = os.path.join(os.path.dirname(__file__), 'scanner.db')
        conn = sqlite3.connect(db_path)
        c = conn.cursor()
        c.execute('SELECT id, network, schedule_type, schedule_value FROM scheduled_tasks WHERE enabled = 1')
        tasks = c.fetchall()
        conn.close()

        formatted_tasks = []
        for task in tasks:
            task_id, network, schedule_type, schedule_value = task
            # 格式化时间显示
            if schedule_type == 'interval':
                time_display = f"每{schedule_value}分钟"
            else:  # cron类型
                time_display = schedule_value

            formatted_tasks.append({
                'id': task_id,
                'network': network,
                'time': time_display,
                'repeat': schedule_type
            })

        return jsonify(formatted_tasks)
    except Exception as e:
        return jsonify({'error': str(e)}), 500

@app.route('/api/schedule', methods=['POST'])
def create_schedule():
    """添加定时任务"""
    try:
        data = request.json
        network = data.get('network')
        time = data.get('time')  # 格式: "HH:MM"
        repeat = data.get('repeat')  # daily, weekly, monthly

        if not all([network, time, repeat]):
            return jsonify({'error': '缺少必要参数'}), 400

        # 解析时间
        hour, minute = map(int, time.split(':'))
        
        # 根据重复类型创建cron表达式
        if repeat == 'daily':
            cron = f"{minute} {hour} * * *"
        elif repeat == 'weekly':
            cron = f"{minute} {hour} * * 0"  # 每周日
        elif repeat == 'monthly':
            cron = f"{minute} {hour} 1 * *"  # 每月1号
        else:
            return jsonify({'error': '无效的重复类型'}), 400

        # 保存到数据库
        db_path = os.path.join(os.path.dirname(__file__), 'scanner.db')
        conn = sqlite3.connect(db_path)
        c = conn.cursor()
        c.execute('''
            INSERT INTO scheduled_tasks (network, schedule_type, schedule_value)
            VALUES (?, ?, ?)
        ''', (network, 'cron', cron))
        task_id = c.lastrowid
        conn.commit()
        conn.close()

        # 添加到调度器
        scheduler.add_job(
            func=scheduled_scan,
            trigger=CronTrigger.from_crontab(cron),
            args=[network, f"scheduled_scan_{task_id}"],
            id=f"task_{task_id}"
        )

        return jsonify({'success': True, 'id': task_id})
    except Exception as e:
        return jsonify({'error': str(e)}), 500

@app.route('/api/schedule/<int:task_id>', methods=['DELETE'])
def delete_schedule(task_id):
    """删除定时任务"""
    try:
        # 从数据库中删除
        db_path = os.path.join(os.path.dirname(__file__), 'scanner.db')
        conn = sqlite3.connect(db_path)
        c = conn.cursor()
        c.execute('UPDATE scheduled_tasks SET enabled = 0 WHERE id = ?', (task_id,))
        conn.commit()
        conn.close()

        # 从调度器中移除
        try:
            scheduler.remove_job(f"task_{task_id}")
        except:
            pass  # 如果任务不存在，忽略错误

        return jsonify({'success': True})
    except Exception as e:
        return jsonify({'error': str(e)}), 500

@app.route('/device_status/<ip>')
def get_device_status(ip):
    """获取设备状态"""
    try:
        db_path = os.path.join(os.path.dirname(__file__), 'scanner.db')
        conn = sqlite3.connect(db_path)
        c = conn.cursor()
        c.execute('SELECT * FROM device_status WHERE ip = ?', (ip,))
        status = c.fetchone()
        conn.close()
        
        if status:
            return jsonify({
                'status': 'success',
                'data': {
                    'ip': status[0],
                    'last_check': status[1],
                    'status': status[2],
                    'response_time': status[3],
                    'error_count': status[4],
                    'last_error': status[5]
                }
            })
        else:
            return jsonify({'status': 'error', 'message': '设备未找到'})
    except Exception as e:
        return jsonify({'status': 'error', 'message': str(e)})

@app.route('/check_status/<ip>')
def check_status(ip):
    """立即检查设备状态"""
    try:
        result = check_device_status(ip)
        return jsonify({'status': 'success', 'data': result})
    except Exception as e:
        return jsonify({'status': 'error', 'message': str(e)})

@app.route('/')
def index():
    return render_template('index.html')

@app.route('/export/<format>')
def export_results(format):
    """导出扫描结果"""
    if not hasattr(app, 'last_scan_results'):
        return jsonify({'error': 'No scan results available'})
    
    timestamp = datetime.datetime.now().strftime('%Y%m%d_%H%M%S')
    
    if format == 'csv':
        output = StringIO()
        writer = csv.writer(output)
        writer.writerow(['IP地址', '设备类型', '响应时间(ms)', '开放端口'])
        
        for result in app.last_scan_results:
            writer.writerow([
                result['ip'],
                result['device_type'],
                result['response_time'],
                ','.join(map(str, result['open_ports']))
            ])
        
        output.seek(0)
        return Response(
            output.getvalue(),
            mimetype='text/csv',
            headers={'Content-Disposition': f'attachment;filename=scan_results_{timestamp}.csv'}
        )
    
    elif format == 'json':
        return Response(
            json.dumps(app.last_scan_results, ensure_ascii=False, indent=2),
            mimetype='application/json',
            headers={'Content-Disposition': f'attachment;filename=scan_results_{timestamp}.json'}
        )
    
    return jsonify({'error': 'Invalid format'})

@app.route('/scan', methods=['POST'])
def scan_network():
    network = request.json.get('network', '')
    try:
        # 验证输入的网络地址格式
        if not network:
            return jsonify({'error': '请输入网络地址'}), 400
            
        # 如果是单个IP地址，转换为CIDR格式
        if '/' not in network:
            network = network + '/32'
            
        # 验证网络地址
        try:
            net = ipaddress.ip_network(network, strict=False)
        except ValueError as e:
            return jsonify({'error': '无效的网络地址格式'}), 400
        
        # 清空之前的结果
        while not scan_results.empty():
            scan_results.get()

        def process_ip_and_queue(ip):
            result = scan_single_ip(str(ip))
            if result:
                scan_results.put(result)

        # 使用线程池进行扫描
        with ThreadPoolExecutor(max_workers=MAX_WORKERS) as executor:
            futures = [executor.submit(process_ip_and_queue, ip) for ip in net.hosts()]
            # 等待所有任务完成
            for future in as_completed(futures):
                try:
                    future.result()  # 获取结果，捕获任何异常
                except Exception as e:
                    print(f"Error processing IP: {str(e)}")
                    continue

        # 收集结果
        results = []
        while not scan_results.empty():
            result = scan_results.get()
            if result:  # 确保结果不为None
                results.append(result)

        # 按IP地址排序
        if results:
            results.sort(key=lambda x: [int(i) for i in x['ip'].split('.')])
            
            try:
                # 保存扫描结果到数据库
                save_scan_results(network, results)
            except Exception as e:
                print(f"Error saving results: {str(e)}")
                # 即使保存失败也继续返回结果
            
            return jsonify({'devices': results})
        else:
            return jsonify({'devices': [], 'message': '未发现活动设备'}), 200
            
    except Exception as e:
        print(f"Scan error: {str(e)}")
        return jsonify({'error': '扫描过程发生错误，请检查网络地址格式是否正确'}), 400

@app.route('/device_history/<ip>')
def get_device_history(ip):
    history = get_device_history(ip)
    if history:
        return jsonify(history)
    return jsonify({'error': 'No history found'})

def try_basic_auth(ip, port, auth_list):
    """尝试基本认证"""
    headers = {
        'Connection': 'close',
        'User-Agent': 'Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/91.0.4472.124 Safari/537.36',
        'Accept': '*/*'
    }
    
    for auth in auth_list:
        try:
            auth_header = base64.b64encode(f"{auth['user']}:{auth['pass']}".encode()).decode()
            headers['Authorization'] = f"Basic {auth_header}"
            
            # 尝试多个可能的API端点
            endpoints = [
                '/ISAPI/System/deviceInfo',
                '/System/deviceInfo',
                '/cgi-bin/magicBox.cgi?action=getSystemInfo',
                '/cgi-bin/system.cgi?action=getDeviceInfo'
            ]
            
            for endpoint in endpoints:
                try:
                    url = f"http://{ip}:{port}{endpoint}"
                    response = requests.get(url, timeout=1, headers=headers, verify=False)
                    
                    if response.status_code == 200 and ("<DeviceInfo>" in response.text or 
                                                      "Hikvision" in response.text or 
                                                      "<firmwareVersion>" in response.text):
                        return response.text, auth, endpoint
                except:
                    continue
                    
        except:
            continue
    return None, None, None

def check_rtsp_auth(ip, auth):
    """检查RTSP认证"""
    try:
        import socket
        s = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
        s.settimeout(1)
        s.connect((ip, 554))
        
        # 发送RTSP OPTIONS请求
        auth_base64 = base64.b64encode(f"{auth['user']}:{auth['pass']}".encode()).decode()
        options_request = (
            f'OPTIONS rtsp://{ip} RTSP/1.0\r\n'
            f'CSeq: 1\r\n'
            f'Authorization: Basic {auth_base64}\r\n'
            f'User-Agent: HIKVISION\r\n\r\n'
        )
        
        s.send(options_request.encode())
        response = s.recv(1024).decode()
        
        if '200 OK' in response:
            return True
            
    except:
        pass
    finally:
        try:
            s.close()
        except:
            pass
    return False

def get_hikvision_info(ip, port=80):
    """获取海康摄像机信息"""
    try:
        # 常用的用户名和密码组合
        auth_list = [
            {'user': 'admin', 'pass': '12345'},
            {'user': 'admin', 'pass': 'admin'},
            {'user': 'admin', 'pass': 'Admin12345'},
            {'user': 'admin', 'pass': ''},
            {'user': 'root', 'pass': '12345'},
            {'user': 'admin', 'pass': 'abcd1234'},
            {'user': 'admin', 'pass': 'hik12345'},
            {'user': 'admin', 'pass': '123456'},
        ]
        
        device_info = {}
        
        # 首先尝试8000端口
        response_text, auth, endpoint = try_basic_auth(ip, 8000, auth_list)
        if response_text:
            device_info['port'] = 8000
            device_info['auth'] = auth
            device_info['endpoint'] = endpoint
            
            # 解析设备信息
            model_match = re.search(r'<model>(.*?)</model>', response_text)
            firmware_match = re.search(r'<firmwareVersion>(.*?)</firmwareVersion>', response_text)
            serial_match = re.search(r'<serialNumber>(.*?)</serialNumber>', response_text)
            
            if model_match:
                device_info['model'] = model_match.group(1)
            if firmware_match:
                device_info['firmware'] = firmware_match.group(1)
            if serial_match:
                device_info['serial'] = serial_match.group(1)
                
            # 检查RTSP可用性
            if check_rtsp_auth(ip, auth):
                device_info['rtsp'] = True
                
            return f"Hikvision {device_info.get('model', 'Camera')} (Firmware: {device_info.get('firmware', 'Unknown')})", device_info

        # 然后尝试80端口
        if port == 80:
            response_text, auth, endpoint = try_basic_auth(ip, 80, auth_list)
            if response_text:
                device_info['port'] = 80
                device_info['auth'] = auth
                device_info['endpoint'] = endpoint
                
                # 解析设备信息
                model_match = re.search(r'<model>(.*?)</model>', response_text)
                firmware_match = re.search(r'<firmwareVersion>(.*?)</firmwareVersion>', response_text)
                serial_match = re.search(r'<serialNumber>(.*?)</serialNumber>', response_text)
                
                if model_match:
                    device_info['model'] = model_match.group(1)
                if firmware_match:
                    device_info['firmware'] = firmware_match.group(1)
                if serial_match:
                    device_info['serial'] = serial_match.group(1)
                    
                # 检查RTSP可用性
                if check_rtsp_auth(ip, auth):
                    device_info['rtsp'] = True
                    
                return f"Hikvision {device_info.get('model', 'Camera')} (Firmware: {device_info.get('firmware', 'Unknown')})", device_info

        # 检查RTSP端口
        try:
            sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
            sock.settimeout(0.5)
            result = sock.connect_ex((str(ip), 554))
            if result == 0:
                device_info['rtsp_port'] = 554
                return "Hikvision Camera (RTSP Available)", device_info
            sock.close()
        except:
            pass

    except Exception as e:
        print(f"Error scanning {ip}: {str(e)}")
    return None, None

def scan_single_ip(ip):
    """扫描单个IP地址"""
    try:
        # 快速ping检测
        response_time = ping(str(ip), timeout=0.5)
        if response_time is not None:
            # 检查关键端口
            open_ports = []
            device_type = None
            key_ports = [80, 8000]  # 优先检查这些端口
            
            for port in key_ports:
                try:
                    sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
                    sock.settimeout(0.5)
                    result = sock.connect_ex((str(ip), port))
                    if result == 0:
                        open_ports.append(port)
                        if not device_type:  # 如果还没有识别出设备类型
                            device_type, device_info = get_hikvision_info(str(ip), port)
                            if device_type:
                                return {
                                    'ip': str(ip),
                                    'status': 'active',
                                    'device_type': device_type,
                                    'response_time': round(response_time * 1000, 2),
                                    'open_ports': sorted(open_ports),
                                    'model': device_info.get('model', ''),
                                    'firmware': device_info.get('firmware', ''),
                                    'auth_info': device_info.get('auth', '')
                                }
                    sock.close()
                except:
                    continue

            # 如果已经确认是Hikvision设备，就不再检查其他端口
            if not device_type:
                # 检查其他端口
                other_ports = [443, 21, 23, 554]  # HTTPS, FTP, Telnet, RTSP
                for port in other_ports:
                    try:
                        sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
                        sock.settimeout(0.3)
                        result = sock.connect_ex((str(ip), port))
                        if result == 0:
                            open_ports.append(port)
                        sock.close()
                    except:
                        continue

            return {
                'ip': str(ip),
                'status': 'active',
                'device_type': device_type or 'Unknown',
                'response_time': round(response_time * 1000, 2),
                'open_ports': sorted(open_ports)
            }
    except:
        pass
    return None

def save_scan_results(network, results):
    """保存扫描结果到数据库"""
    try:
        db_path = os.path.join(os.path.dirname(__file__), 'scanner.db')
        conn = sqlite3.connect(db_path)
        c = conn.cursor()
        
        # 将结果转换为JSON字符串
        results_json = json.dumps(results)
        
        # 插入扫描历史
        c.execute('INSERT INTO scan_history (scan_time, network, results) VALUES (?, ?, ?)',
                  (datetime.datetime.now(), network, results_json))
        
        # 更新设备状态表
        for device in results:
            c.execute('''
                INSERT OR REPLACE INTO device_status 
                (ip, last_check, status, response_time, device_type, open_ports)
                VALUES (?, ?, ?, ?, ?, ?)
            ''', (
                device['ip'],
                datetime.datetime.now(),
                device['status'],
                device['response_time'],
                device.get('device_type', 'Unknown'),
                ','.join(map(str, device.get('open_ports', [])))
            ))
        
        conn.commit()
        conn.close()
        return True
    except Exception as e:
        print(f"Error saving results: {str(e)}")
        return False

def get_device_history(ip):
    db_path = os.path.join(os.path.dirname(__file__), 'scanner.db')
    conn = sqlite3.connect(db_path)
    c = conn.cursor()
    c.execute('SELECT * FROM device_info WHERE ip = ?', (ip,))
    result = c.fetchone()
    conn.close()
    if result:
        return {
            'ip': result[0],
            'last_seen': result[1],
            'device_type': result[2],
            'model': result[3],
            'firmware': result[4],
            'ports': json.loads(result[5]),
            'auth_info': result[6]
        }
    return None

if __name__ == '__main__':
    init_db()
    load_scheduled_tasks()
    app.run(host='0.0.0.0', port=5000, debug=True, threaded=True)
